use anyhow::Result;
use serde::{Deserialize, Serialize};
use shared_memory::*;
use std::io::Cursor;
use zmq::{Context as ZmqContext, REP};

// 图像处理函数 (与之前相同)
mod image_processing {
    use image::imageops;
    use image::{DynamicImage, GenericImageView, RgbaImage};

    /// 灰度处理
    pub fn process_grayscale(img: &DynamicImage) -> RgbaImage {
        img.grayscale().to_rgba8()
    }

    /// 颜色反转
    pub fn process_invert(img: &DynamicImage) -> RgbaImage {
        let mut img_rgba = img.to_rgba8();
        imageops::invert(&mut img_rgba);
        img_rgba
    }

    /// 高斯模糊
    pub fn process_blur(img: &DynamicImage) -> RgbaImage {
        img.blur(5.0).to_rgba8()
    }

    /// 旋转90度
    pub fn process_rotate90(img: &DynamicImage) -> RgbaImage {
        img.rotate90().to_rgba8()
    }

    /// Sobel边缘检测
    pub fn process_canny_edge(img: &DynamicImage) -> RgbaImage {
        let gray_img = img.to_luma8();
        let mut edges = gray_img.clone();

        for y in 1..gray_img.height() - 1 {
            for x in 1..gray_img.width() - 1 {
                let mut gx = 0i16;
                let mut gy = 0i16;

                for dy in -1..=1 {
                    for dx in -1..=1 {
                        let pixel = gray_img
                            .get_pixel((x as i32 + dx) as u32, (y as i32 + dy) as u32)[0]
                            as i16;

                        match dx {
                            -1 => gx -= pixel,
                            1 => gx += pixel,
                            _ => {}
                        }

                        match dy {
                            -1 => gy -= pixel,
                            1 => gy += pixel,
                            _ => {}
                        }
                    }
                }

                let magnitude = (gx.abs() + gy.abs()).min(255) as u8;
                edges.put_pixel(x, y, image::Luma([magnitude]));
            }
        }

        DynamicImage::ImageLuma8(edges).to_rgba8()
    }
}

#[derive(Serialize, Deserialize, Debug)]
struct ControlMsg {
    command: String,
    input_size: usize, // 输入数据大小
    shmem_size: usize, // 共享内存总大小
    params: Option<serde_json::Value>,
}

#[derive(Serialize, Deserialize)]
struct ResponseMsg {
    status: String,
    output_size: usize, // 输出数据大小
    error: Option<String>,
}

/// 处理图像请求
fn process_image(input: &[u8], command: &str) -> Result<Vec<u8>> {
    // 加载图像
    let img = image::load_from_memory(input)?;

    // 根据命令选择处理函数
    let result = match command {
        "grayscale" => image_processing::process_grayscale(&img),
        "invert" => image_processing::process_invert(&img),
        "blur" => image_processing::process_blur(&img),
        "rotate90" => image_processing::process_rotate90(&img),
        "canny_edge" => image_processing::process_canny_edge(&img),
        _ => return Err(anyhow::anyhow!("Unknown command: {}", command)),
    };

    // 编码为PNG
    let mut buffer = Vec::new();
    result.write_to(&mut Cursor::new(&mut buffer), image::ImageOutputFormat::Png)?;

    Ok(buffer)
}

fn main() -> Result<()> {
    let ctx = ZmqContext::new();
    let socket = ctx.socket(REP)?;

    // 跨平台IPC路径
    #[cfg(not(target_os = "windows"))]
    let ipc_path = "ipc:///tmp/image_processor.ipc";
    #[cfg(target_os = "windows")]
    let ipc_path = "tcp://127.0.0.1:5555";

    socket.bind(ipc_path)?;
    println!("Rust processor started. Listening on: {}", ipc_path);

    // 访问共享内存
    let shmem = match ShmemConf::new().os_id(&"pyimg").open() {
        Ok(m) => m,
        Err(e) => {
            let response = ResponseMsg {
                status: "error".to_string(),
                output_size: 0,
                error: Some(format!("Shmem open failed: {}", e)),
            };
            socket.send(&serde_json::to_string(&response)?, 0)?;
            return Err(e.into());
        }
    };

    loop {
        let msg_str = match socket.recv_string(0) {
            Ok(Ok(msg)) => msg,
            Ok(Err(e)) => {
                eprintln!("Invalid UTF-8 message: {:?}", e);
                continue;
            }
            Err(e) => {
                eprintln!("Error receiving message: {}", e);
                continue;
            }
        };

        let control: ControlMsg = match serde_json::from_str(&msg_str) {
            Ok(c) => c,
            Err(e) => {
                let response = ResponseMsg {
                    status: "error".to_string(),
                    output_size: 0,
                    error: Some(format!("Invalid control message: {}", e)),
                };
                socket.send(&serde_json::to_string(&response)?, 0)?;
                continue;
            }
        };

        // 验证内存大小
        if control.input_size > control.shmem_size {
            let response = ResponseMsg {
                status: "error".to_string(),
                output_size: 0,
                error: Some(format!(
                    "Input size {} exceeds shared memory size {}",
                    control.input_size, control.shmem_size
                )),
            };
            socket.send(&serde_json::to_string(&response)?, 0)?;
            continue;
        }

        // 从共享内存读取输入数据
        let input_data = unsafe { std::slice::from_raw_parts(shmem.as_ptr(), control.input_size) };

        // 处理图像
        let result = process_image(input_data, &control.command);

        match result {
            Ok(output_data) => {
                // 验证输出数据是否适合共享内存
                if output_data.len() > (control.shmem_size - control.input_size) {
                    let response = ResponseMsg {
                        status: "error".to_string(),
                        output_size: 0,
                        error: Some(format!(
                            "Output size {} exceeds available space {}",
                            output_data.len(),
                            control.shmem_size - control.input_size
                        )),
                    };
                    socket.send(&serde_json::to_string(&response)?, 0)?;
                    continue;
                }

                // 将输出数据写入共享内存（从输入数据之后开始）
                let output_ptr = unsafe { shmem.as_ptr().add(control.input_size) };
                unsafe {
                    std::ptr::copy_nonoverlapping(
                        output_data.as_ptr(),
                        output_ptr,
                        output_data.len(),
                    );
                }

                let response = ResponseMsg {
                    status: "success".to_string(),
                    output_size: output_data.len(),
                    error: None,
                };
                socket.send(&serde_json::to_string(&response)?, 0)?;
            }
            Err(e) => {
                let response = ResponseMsg {
                    status: "error".to_string(),
                    output_size: 0,
                    error: Some(e.to_string()),
                };
                socket.send(&serde_json::to_string(&response)?, 0)?;
            }
        }
    }
}
